package cc.decat.escript.ir;

import cc.decat.escript.codegen.BlockGenerate;
import cc.decat.escript.codegen.Codegen;
import cc.decat.escript.codegen.FunctionGenerate;
import cc.decat.escript.ir.expression.*;
import cc.decat.escript.ir.statement.*;
import cc.decat.escript.runtime.Opcode;

import java.util.*;

public class CodegenVisitorImpl implements Visitor {
    private static class LoopState {
        List<Integer> breakInstructionList;
        List<Integer> incInstructionList;
        public LoopState() {
            this.breakInstructionList = new ArrayList<>();
            this.incInstructionList = new ArrayList<>();
        }
    }

    private Deque<LoopState> loopStateDeque = new ArrayDeque<>();

    private Codegen codegen;

    public CodegenVisitorImpl(Codegen codegen) {
        this.codegen = codegen;
    }
    @Override
    public void visit(ChunkNode chunkNode) {
        codegen.generateChunk(chunkNode);
        if(chunkNode.getStatementList() != null) {
            chunkNode.getStatementList().accept(this);
        }
    }

    @Override
    public void visit(VariableExpression variableExpression) {
        codegen.loadVar(variableExpression.getName());
    }

    @Override
    public void visit(LiteralExpression literalExpression) {
        codegen.loadConst(literalExpression);
    }

    @Override
    public void visit(AssignStatement assignStatement) {
        for(Node node : assignStatement.getExpressionList()) {
            node.accept(this);
        }
        Collections.reverse(assignStatement.getAssignableList());
        for(Node node : assignStatement.getAssignableList()) {
            AssignableNode assignableNode = (AssignableNode) node;
            switch(assignableNode.getType()) {
                case Variable: {
                    codegen.storeVar(assignableNode.getVariableName());
                    break;
                }
                case ListIndex: {
                    assignableNode.getListExpression().accept(this);
                    assignableNode.getIndexExpression().accept(this);
                    codegen.addInstruction(Opcode.STORE_SUBSCR);
                    break;
                }
                case MemberAttribute: {
                    assignableNode.getObjectExpression().accept(this);
                    codegen.loadConst(new LiteralExpression().setStringValue(assignableNode.getAttributeName()));
                    codegen.addInstruction(Opcode.STORE_ATTR);
                    break;
                }
                default: break;
            }
        }
    }

    @Override
    public void visit(StatementList statementList) {
        for(Node statement : statementList.getStatementList()) {
            statement.accept(this);
        }
    }

    @Override
    public void visit(BlockStatement blockStatement) {
        // entry block
        codegen.entryBlock();
        blockStatement.getStatementList().accept(this);
        // leave block
        codegen.leaveBlock();
    }

    @Override
    public void visit(BinaryExpression binaryExpression) {
        if(binaryExpression.getOpType() == BinaryExpression.OpType.LOGIC_AND) {
            binaryExpression.getLeft().accept(this);
            int mark = codegen.mark();
            codegen.addInstruction(Opcode.JUMP_IF_FALSE_OR_POP);
            binaryExpression.getRight().accept(this);
            codegen.getInstruction(mark).setOperand(codegen.mark());
            return;
        } else if(binaryExpression.getOpType() == BinaryExpression.OpType.LOGIC_OR) {
            binaryExpression.getLeft().accept(this);
            int mark = codegen.mark();
            codegen.addInstruction(Opcode.JUMP_IF_TRUE_OR_POP);
            binaryExpression.getRight().accept(this);
            codegen.getInstruction(mark).setOperand(codegen.mark());
            return;
        }
        binaryExpression.getLeft().accept(this);
        binaryExpression.getRight().accept(this);
        switch(binaryExpression.getOpType()) {
            case ADD: {
                codegen.addInstruction(Opcode.BINARY_ADD);
                break;
            }
            case SUB: {
                codegen.addInstruction(Opcode.BINARY_SUB);
                break;
            }
            case MUL: {
                codegen.addInstruction(Opcode.BINARY_MUL);
                break;
            }
            case DIV: {
                codegen.addInstruction(Opcode.BINARY_DIV);
                break;
            }
            case MOD: {
                codegen.addInstruction(Opcode.BINARY_MOD);
                break;
            }
            case LT: {
                codegen.addInstruction(Opcode.COMPARE_LT);
                break;
            }
            case LTE: {
                codegen.addInstruction(Opcode.COMPARE_LTE);
                break;
            }
            case GT: {
                codegen.addInstruction(Opcode.COMPARE_GT);
                break;
            }
            case GTE: {
                codegen.addInstruction(Opcode.COMPARE_GTE);
                break;
            }
            case EQ: {
                codegen.addInstruction(Opcode.COMPARE_EQ);
                break;
            }
            case NEQ: {
                codegen.addInstruction(Opcode.COMPARE_NEQ);
                break;
            }
            default: break;
        }
    }

    @Override
    public void visit(CallFunctionStatement callFunctionStatement) {
        for(Node arg : callFunctionStatement.getArgs()) {
            arg.accept(this);
        }
        callFunctionStatement.getExpression().accept(this);
        codegen.addInstruction(Opcode.CALL_FUNCTION, callFunctionStatement.getArgs().size());
        codegen.addInstruction(Opcode.REMOVE_TOP);
    }

    @Override
    public void visit(FunctionDeclarationStatement functionDeclarationStatement) {
        codegen.generateEntryFunction(functionDeclarationStatement);
        ((BlockStatement) functionDeclarationStatement.getBlock()).getStatementList().accept(this);
        codegen.generateLeaveFunction();
    }

    @Override
    public void visit(IfStatement ifStatement) {
        List<Integer> quickIfMarkList = new ArrayList<>();

        int jumpElseMark = -1;
        int jumpElseIfMark = -1;

        for(int i = 0; i < ifStatement.getCondExpr().size(); i++) {
            if(jumpElseIfMark != -1) {
                codegen.getInstruction(jumpElseIfMark).setOperand(codegen.mark());
            }
            ifStatement.getCondExpr().get(i).accept(this);
            jumpElseIfMark = codegen.mark();
            codegen.addInstruction(Opcode.JUMP_IF_FALSE);
            codegen.setNextBlockType(BlockGenerate.Type.IF);
            ifStatement.getThenBody().get(i).accept(this);
            quickIfMarkList.add(codegen.mark());
            codegen.addInstruction(Opcode.JUMP);
            if(i == ifStatement.getCondExpr().size() - 1) {
                jumpElseMark = jumpElseIfMark;
            }
        }

        codegen.getInstruction(jumpElseMark).setOperand(codegen.mark());
        if(ifStatement.getElseBody() != null) {
            codegen.setNextBlockType(BlockGenerate.Type.IF);
            ifStatement.getElseBody().accept(this);
        }

        for(Integer index : quickIfMarkList) {
            codegen.getInstruction(index).setOperand(codegen.mark());
        }
        codegen.addInstruction(Opcode.NOP);
    }

    @Override
    public void visit(CallFunctionExpression callFunctionExpression) {
        for(Node arg : callFunctionExpression.getArgs()) {
            arg.accept(this);
        }
        callFunctionExpression.getExpression().accept(this);
        codegen.addInstruction(Opcode.CALL_FUNCTION, callFunctionExpression.getArgs().size());
    }

    @Override
    public void visit(ReturnStatement returnStatement) {
        if(returnStatement.getExpression() != null) {
            returnStatement.getExpression().accept(this);
        } else {
            codegen.addInstruction(Opcode.LOAD_NULL);
        }
        codegen.addInstruction(Opcode.RETURN_VALUE);
    }

    @Override
    public void visit(AnonymousFunctionExpression anonymousFunctionExpression) {
        String lambdaName = codegen.generateEntryLambda(anonymousFunctionExpression);
        ((BlockStatement) anonymousFunctionExpression.getBlock()).getStatementList().accept(this);
        codegen.generateLeaveLambda();
        codegen.addInstruction(Opcode.LOAD_LOCAL_VAR, codegen.getBlockGenerate().getSymbolTable().findVar(lambdaName).getIndex());
    }

    @Override
    public void visit(ForStatement forStatement) {
        // 此block用于隔离for循环的作用域，使在initialize和increment定义变量不会影响到外部作用域
        codegen.entryBlock();

        if(forStatement.getInitialization() != null) {
            forStatement.getInitialization().accept(this);
        }

        // 记录循环开始标记，完成一次循环后会再次jump到这个地方
        int startLoopMark = codegen.mark();
        loopStateDeque.addLast(new LoopState());

        // 生成循环条件，如果条件为空的话则是无限循环
        if(forStatement.getCondition() != null) {
            forStatement.getCondition().accept(this);
        } else {
            codegen.addInstruction(Opcode.LOAD_TRUE);
        }

        // 记录循环退出标记
        loopStateDeque.getLast().breakInstructionList.add(codegen.mark());
        codegen.addInstruction(Opcode.JUMP_IF_FALSE);

        // 生成循环体
        codegen.setNextBlockType(BlockGenerate.Type.LOOP);
        forStatement.getBody().accept(this);

        // 记录increment标记
        int incMark = codegen.mark();
        if(forStatement.getIncrement() != null) {
            forStatement.getIncrement().accept(this);
        }
        codegen.addInstruction(Opcode.JUMP, startLoopMark);

        // 循环结束标记
        int quitMark = codegen.mark();
        codegen.leaveBlock();

        for(Integer index : loopStateDeque.getLast().incInstructionList) {
            codegen.getInstruction(index).setOperand(incMark);
        }

        for(Integer index : loopStateDeque.getLast().breakInstructionList) {
            codegen.getInstruction(index).setOperand(quitMark);
        }

        loopStateDeque.removeLast();
    }

    @Override
    public void visit(ListExpression listExpression) {
        for(Node expr : listExpression.getList()) {
            expr.accept(this);
        }
        codegen.addInstruction(Opcode.BUILD_LIST, listExpression.getList().size());
    }

    @Override
    public void visit(ListIndexExpression listIndexExpression) {
        listIndexExpression.getExpression().accept(this);
        listIndexExpression.getIndex().accept(this);
        codegen.addInstruction(Opcode.LOAD_SUBSCR);
    }

    @Override
    public void visit(DictExpression dictExpression) {
        for(Node key : dictExpression.getMap().keySet()) {
            key.accept(this);
            dictExpression.getMap().get(key).accept(this);
        }
        codegen.addInstruction(Opcode.BUILD_DICT, dictExpression.getMap().keySet().size());
    }

    @Override
    public void visit(ContinueStatement continueStatement) {
        if(loopStateDeque.isEmpty()) {
            System.err.println("continue指令不在循环块中");
            return;
        }
        loopStateDeque.getLast().incInstructionList.add(codegen.mark());
        codegen.addInstruction(Opcode.CONTINUE);
    }

    @Override
    public void visit(BreakStatement breakStatement) {
        if(loopStateDeque.isEmpty()) {
            System.err.println("break指令不在循环块中");
            return;
        }
        loopStateDeque.getLast().breakInstructionList.add(codegen.mark());
        codegen.addInstruction(Opcode.BREAK);
    }

    @Override
    public void visit(MemberExpression memberExpression) {
        memberExpression.getExpression().accept(this);
        codegen.loadConst(new LiteralExpression().setStringValue(memberExpression.getName()));
        codegen.addInstruction(Opcode.LOAD_ATTR);
    }
}
